"""
API网关负载均衡服务
==================

基于RedFire框架的负载均衡服务
负责服务实例管理和负载均衡算法
"""

import logging
import random
import time
from typing import Dict, Any, Optional, List
from collections import defaultdict

from redfire_framework.cache.service import CacheService
from redfire_framework.events.bus import EventBus

logger = logging.getLogger(__name__)


class LoadBalancerService:
    """
    负载均衡服务
    
    负责：
    - 服务实例管理
    - 负载均衡算法
    - 健康状态跟踪
    """
    
    def __init__(
        self,
        cache_service: Optional[CacheService] = None,
        event_bus: Optional[EventBus] = None,
        config: Optional[Dict[str, Any]] = None
    ):
        """
        初始化负载均衡服务
        
        Args:
            cache_service: 缓存服务
            event_bus: 事件总线
            config: 配置信息
        """
        self.cache_service = cache_service
        self.event_bus = event_bus
        self.config = config or {}
        
        # 服务实例配置
        self.services = self.config.get("services", {})
        
        # 负载均衡配置
        self.lb_config = self.config.get("load_balancing", {})
        self.strategy = self.lb_config.get("strategy", "weighted_round_robin")
        
        # 轮询计数器
        self._round_robin_counters = defaultdict(int)
        
        logger.info(f"Load balancer service initialized with strategy: {self.strategy}")
    
    async def list_services(self) -> Dict[str, List[Dict[str, Any]]]:
        """
        列出所有服务实例
        
        Returns:
            Dict[str, List[Dict[str, Any]]]: 服务实例列表
        """
        result = {}
        for service_name, instances in self.services.items():
            result[service_name] = [
                {
                    "host": instance["host"],
                    "port": instance["port"],
                    "weight": instance.get("weight", 100),
                    "health_check_path": instance.get("health_check_path", "/health"),
                    "timeout": instance.get("timeout", 30.0),
                    "metadata": instance.get("metadata", {})
                }
                for instance in instances
            ]
        return result
    
    async def list_service_instances(self, service_name: str) -> List[Dict[str, Any]]:
        """
        列出指定服务的实例
        
        Args:
            service_name: 服务名称
            
        Returns:
            List[Dict[str, Any]]: 服务实例列表
        """
        instances = self.services.get(service_name, [])
        return [
            {
                "host": instance["host"],
                "port": instance["port"],
                "weight": instance.get("weight", 100),
                "health_check_path": instance.get("health_check_path", "/health"),
                "timeout": instance.get("timeout", 30.0),
                "metadata": instance.get("metadata", {}),
                "healthy": await self._is_instance_healthy(service_name, instance)
            }
            for instance in instances
        ]
    
    async def get_service_instance(self, service_name: str) -> Optional[Dict[str, Any]]:
        """
        根据负载均衡策略获取服务实例
        
        Args:
            service_name: 服务名称
            
        Returns:
            Optional[Dict[str, Any]]: 选中的服务实例
        """
        instances = self.services.get(service_name, [])
        if not instances:
            return None
        
        # 过滤健康的实例
        healthy_instances = []
        for instance in instances:
            if await self._is_instance_healthy(service_name, instance):
                healthy_instances.append(instance)
        
        if not healthy_instances:
            logger.warning(f"No healthy instances for service: {service_name}")
            # 如果没有健康实例，返回第一个实例作为fallback
            return instances[0] if instances else None
        
        # 根据策略选择实例
        if self.strategy == "round_robin":
            return self._round_robin_select(service_name, healthy_instances)
        elif self.strategy == "weighted_round_robin":
            return self._weighted_round_robin_select(service_name, healthy_instances)
        elif self.strategy == "random":
            return self._random_select(healthy_instances)
        elif self.strategy == "least_connections":
            return await self._least_connections_select(service_name, healthy_instances)
        else:
            # 默认使用轮询
            return self._round_robin_select(service_name, healthy_instances)
    
    def _round_robin_select(self, service_name: str, instances: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        轮询选择
        
        Args:
            service_name: 服务名称
            instances: 实例列表
            
        Returns:
            Dict[str, Any]: 选中的实例
        """
        counter = self._round_robin_counters[service_name]
        selected = instances[counter % len(instances)]
        self._round_robin_counters[service_name] = (counter + 1) % len(instances)
        return selected
    
    def _weighted_round_robin_select(self, service_name: str, instances: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        加权轮询选择
        
        Args:
            service_name: 服务名称
            instances: 实例列表
            
        Returns:
            Dict[str, Any]: 选中的实例
        """
        # 构建加权列表
        weighted_instances = []
        for instance in instances:
            weight = instance.get("weight", 100)
            weighted_instances.extend([instance] * weight)
        
        if not weighted_instances:
            return instances[0]
        
        counter = self._round_robin_counters[service_name]
        selected = weighted_instances[counter % len(weighted_instances)]
        self._round_robin_counters[service_name] = (counter + 1) % len(weighted_instances)
        return selected
    
    def _random_select(self, instances: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        随机选择
        
        Args:
            instances: 实例列表
            
        Returns:
            Dict[str, Any]: 选中的实例
        """
        return random.choice(instances)
    
    async def _least_connections_select(self, service_name: str, instances: List[Dict[str, Any]]) -> Dict[str, Any]:
        """
        最少连接选择
        
        Args:
            service_name: 服务名称
            instances: 实例列表
            
        Returns:
            Dict[str, Any]: 选中的实例
        """
        # 获取每个实例的连接数
        instance_connections = {}
        for instance in instances:
            instance_key = f"{instance['host']}:{instance['port']}"
            connections = 0
            
            if self.cache_service:
                connections = await self.cache_service.get(f"connections:{service_name}:{instance_key}") or 0
            
            instance_connections[instance_key] = connections
        
        # 选择连接数最少的实例
        min_connections = min(instance_connections.values())
        candidates = [
            instance for instance in instances
            if instance_connections[f"{instance['host']}:{instance['port']}"] == min_connections
        ]
        
        # 如果有多个候选，随机选择
        return random.choice(candidates)
    
    async def _is_instance_healthy(self, service_name: str, instance: Dict[str, Any]) -> bool:
        """
        检查实例是否健康
        
        Args:
            service_name: 服务名称
            instance: 实例信息
            
        Returns:
            bool: 是否健康
        """
        if not self.cache_service:
            return True  # 没有缓存服务，假设健康
        
        instance_key = f"{instance['host']}:{instance['port']}"
        health_key = f"health:{service_name}:{instance_key}"
        
        # 从缓存获取健康状态
        health_status = await self.cache_service.get(health_key)
        if health_status is None:
            return True  # 没有健康状态信息，假设健康
        
        return health_status.get("healthy", True)
    
    async def update_instance_health(self, service_name: str, instance: Dict[str, Any], healthy: bool):
        """
        更新实例健康状态
        
        Args:
            service_name: 服务名称
            instance: 实例信息
            healthy: 是否健康
        """
        if not self.cache_service:
            return
        
        instance_key = f"{instance['host']}:{instance['port']}"
        health_key = f"health:{service_name}:{instance_key}"
        
        health_status = {
            "healthy": healthy,
            "timestamp": time.time(),
            "service_name": service_name,
            "instance": instance_key
        }
        
        await self.cache_service.set(health_key, health_status, ttl=300)  # 5分钟缓存
        
        # 发布健康状态变更事件
        if self.event_bus:
            await self.event_bus.publish("gateway.instance.health_changed", {
                "service_name": service_name,
                "instance": instance_key,
                "healthy": healthy,
                "timestamp": health_status["timestamp"]
            })
    
    async def record_connection(self, service_name: str, instance: Dict[str, Any]):
        """
        记录连接
        
        Args:
            service_name: 服务名称
            instance: 实例信息
        """
        if not self.cache_service:
            return
        
        instance_key = f"{instance['host']}:{instance['port']}"
        connections_key = f"connections:{service_name}:{instance_key}"
        
        current_connections = await self.cache_service.get(connections_key) or 0
        await self.cache_service.set(connections_key, current_connections + 1, ttl=60)
    
    async def release_connection(self, service_name: str, instance: Dict[str, Any]):
        """
        释放连接
        
        Args:
            service_name: 服务名称
            instance: 实例信息
        """
        if not self.cache_service:
            return
        
        instance_key = f"{instance['host']}:{instance['port']}"
        connections_key = f"connections:{service_name}:{instance_key}"
        
        current_connections = await self.cache_service.get(connections_key) or 0
        if current_connections > 0:
            await self.cache_service.set(connections_key, current_connections - 1, ttl=60)
    
    async def refresh_services(self) -> Dict[str, Any]:
        """
        刷新服务实例状态
        
        Returns:
            Dict[str, Any]: 刷新结果
        """
        try:
            # 重新加载服务配置
            self.services = self.config.get("services", {})
            
            # 清除健康状态缓存
            if self.cache_service:
                # 这里应该清除所有健康状态缓存
                # 简化实现
                pass
            
            # 发布服务刷新事件
            if self.event_bus:
                await self.event_bus.publish("gateway.services.refreshed", {
                    "services_count": len(self.services),
                    "timestamp": time.time()
                })
            
            logger.info(f"Services refreshed: {len(self.services)} services")
            
            return {
                "status": "success",
                "services_count": len(self.services)
            }
            
        except Exception as e:
            logger.error(f"Failed to refresh services: {e}", exc_info=True)
            return {
                "status": "error",
                "error": str(e)
            }
